import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from gurobipy import *


class Node:
    def __init__(self):
        self.model = None # copy gurobi里面的LP模型
        self.is_integer = False
        self.model_solution = 0
        self.is_pruned = False # 被剪枝状态
        # 先分这几个试试看吧

    def DeepCopyNode(node):
        new_node = Node()
        new_node.model = node.model.copy() # 把节点的LP模型复制下来
        new_node.is_integer = node.is_integer
        new_node.is_pruned = node.is_pruned
        return new_node

def whether_integer(node):
    i = 0
    for var in node.model.getVars():
        i = i+1
        if(int(var.x) != var.x):
            node.is_integer = False
            break
        elif(i == len(node.model.getVars())):
            node.is_integer = True




if __name__ == '__main__':

    c = [1,1,2]
    a = [[7,2,3],[5,4,7],[2,3,5]]
    b = [36,42,28]

    initial_LP = gurobipy.Model()
    x = {}

    num = len(c)

    for i in range(num):
        x[i] = initial_LP.addVar(lb = 0, ub = GRB.INFINITY , vtype = GRB.CONTINUOUS , name = 'x_' + str(i))
    
    obj = LinExpr(0)
    for i in range(num):  
        obj.addTerms(c[i],x[i])
    initial_LP.setObjective(obj,GRB.MAXIMIZE)

    for i in range(num):
        expr = LinExpr(0)
        for j in range(num):
            expr.addTerms(a[i][j],x[j])
        initial_LP.addConstr(expr <= b[i])

    initial_LP.optimize()

    for var in initial_LP.getVars():
        print(var.Varname, ' = ' , var.x)
    initial_LP.Status


    globle_LB = 0
    globle_UB = np.inf
    eps = 0.001

    optimal_solution = {}



    Queue = []
    node = Node()
    node.model_solution = initial_LP.ObjVal
    node.model = initial_LP.copy()
    depth = 0
    Queue.append(node)
    while( len(Queue) > 0 ):
        depth = depth + 1
        current_node = Queue[0]
        Queue.remove(current_node)
        current_node.model.optimize()
        solution_state = current_node.model.status
        # 判断这个解是否可行
        if(solution_state == 2):
            whether_integer(current_node)
            if (current_node.is_integer == True):
                x_list = []
                print('==============整数解=================')
                for var in current_node.model.getVars():
                    print(var.Varname, ' = ' , var.x)
                    x_list.append(var.x)
                print('=====================================')
                current_node.model_solution = current_node.model.ObjVal
                if (current_node.model_solution != globle_LB):
                    optimal_solution[current_node.model_solution] = []
                if (current_node.model_solution >= globle_LB):
                    # 得到整数解之后先判断，判断之后更新下界
                    globle_LB = current_node.model_solution
                    # print('globle_LB = ',globle_LB)
                    # print('globle_UB = ',globle_UB)
                    optimal_solution[current_node.model_solution].append(x_list)

            else:
                # print('不是整数解')
                current_node.model_solution = current_node.model.ObjVal
                if (current_node.model_solution < globle_LB or current_node.model_solution > globle_UB):
                    current_node.is_pruned = True
                else:
                    globle_UB = current_node.model_solution
                    # 选值最大的变量分支
                    max_x_value = -np.inf
                    for var in current_node.model.getVars():
                        if(int(var.x) != var.x):
                            if(var.x > max_x_value):
                                max_x_value = var.x
                                max_x_name =  var.VarName
                    # 分支
                    left_node = Node.DeepCopyNode(current_node)
                    temp_var = left_node.model.getVarByName(max_x_name)
                    left_node.model.addConstr(temp_var <= int(max_x_value))
                    left_node.model.update()
                    Queue.append(left_node)

                    right_node = Node.DeepCopyNode(current_node)
                    temp_var = right_node.model.getVarByName(max_x_name)
                    right_node.model.addConstr(temp_var >= int(max_x_value) + 1)
                    right_node.model.update()
                    Queue.append(right_node)

        else:
            current_node.is_pruned = True

        # break


    print('迭代结束')
    print('optimal_solution = ',globle_LB)

    print(optimal_solution)
    cnt = 1
    for solution in optimal_solution[globle_LB]:
        print('------第{}组最优解------'.format(cnt))
        print(solution)
        cnt = cnt + 1
    print('-----------------------')

    

